import csv
import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os

from torchvision import transforms


class SequenceDataset(Dataset):
    def __init__(self, csv_file, transform=None, sequence_length=40, split_ratio=0.5,image_prefix_path=None):
        """
        Args:
            csv_file (str): CSV文件路径，每行包含sequence_length个图片路径
            transform (callable, optional): 应用于图片的变换
            sequence_length (int): 每行的图片数量（默认40）
            split_ratio (float): 前split_ratio部分作为输入，后(1-split_ratio)作为目标（默认0.5）
        """
        self.transform = transform

        self.image_prefix_path = image_prefix_path

        self.sequences = []
        self.sequences_input_image_tensor = []
        self.sequences_target_image_tensor = []

        # 读取CSV文件
        with open(csv_file, 'r') as f:
            reader = csv.reader(f)
            for row in reader:
                if len(row) != sequence_length:
                    print(f"Warning: Skipping invalid row with {len(row)} elements")
                    continue
                full_sequence = [os.path.expanduser(path.strip()) for path in row]
                self.sequences.append(full_sequence)

                # self.sequences_input_image_tensor.append(self.load_image(os.path.join(self.image_prefix_path,f"radar_{path}")) for path in full_sequence[:self.split_idx])
                # self.sequences_target_image_tensor.append(self.load_image(os.path.join(self.image_prefix_path,f"radar_{path}")) for path in full_sequence[self.split_idx:])



        self.split_idx = int(sequence_length * split_ratio)

    # 加载图片并应用变换
    def load_image(self,path):
        try:
            img = Image.open(path).convert('L')
            if self.transform:
                img = self.transform(img)
            return img
        except Exception as e:
            print(f"Error loading image {path}: {str(e)}")
            return torch.zeros(1, 560, 480)  # 返回空白图片作为占位符

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        # 获取完整序列
        full_sequence = self.sequences[idx]

        # 分割为输入和目标
        input_sequence = full_sequence[:self.split_idx]
        target_sequence = full_sequence[self.split_idx:]

        input_images = [self.load_image(os.path.join(self.image_prefix_path,f"radar_{path}")) for path in input_sequence]
        target_images = [self.load_image(os.path.join(self.image_prefix_path,f"radar_{path}")) for path in target_sequence]

        # 堆叠为张量 [sequence_length, channels, height, width]
        input_tensor = torch.stack(input_images)
        target_tensor = torch.stack(target_images)


        # input_tensor_shape = input_tensor.shape
        # input_tensor =input_tensor.reshape(input_tensor_shape[1],input_tensor_shape[0],input_tensor_shape[3],input_tensor_shape[2])
        # target_tensor_shape = target_tensor.shape
        # target_tensor = target_tensor.reshape(target_tensor_shape[1], target_tensor_shape[0], target_tensor_shape[3],
        #                                   target_tensor_shape[2])

        input_tensor = input_tensor.permute(1, 0, 3, 2)
        target_tensor = target_tensor.permute(1, 0, 3, 2)




        return input_tensor, target_tensor

if __name__ == '__main__':


    # 定义数据变换（可选）
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    ])

    # 创建Dataset和DataLoader
    dataset = SequenceDataset(
        csv_file=r'D:\python_code\weather\data\TestA\TestA.csv',
        image_prefix_path=r'D:\python_code\weather\data\TestA\TestA\Radar',
        transform=transform,
        sequence_length=40,
        split_ratio=0.5  # 前20张作为输入，后20张作为目标
    )

    dataloader = DataLoader(
        dataset,
        batch_size=4,
        shuffle=True,
        num_workers=4,
        pin_memory=True  # 如果使用GPU，建议开启
    )

    # 使用示例
    for batch_idx, (inputs, targets) in enumerate(dataloader):
        print(f"Batch {batch_idx}:")
        print(f"Input shape: {inputs.shape}")  # [batch_size, 20, channels, height, width]
        print(f"Target shape: {targets.shape}")  # [batch_size, 20, channels, height, width]

        # 如果使用GPU
        # inputs = inputs.to(device)
        # targets = targets.to(device)